#include <iostream>
#include <cstring>
#include <unistd.h>
#include <arpa/inet.h>
#include <thread>
#include <atomic>
#include <vector>
#include <string>
#include <mutex>
#include <algorithm>  // 添加这个头文件

class BridgeServer {
private:
    int server_fd;
    int port;
    std::atomic<bool> running;
    std::thread serverThread;
    std::vector<int> clientSockets;
    std::mutex clientMutex;  // 线程安全保护

public:
    BridgeServer(int port) : port(port), running(false) {}
    ~BridgeServer() { stop(); }

    bool start() {
        // 创建套接字
        server_fd = socket(AF_INET, SOCK_STREAM, 0);
        if (server_fd == -1) {
            std::cerr << "Failed to create socket" << std::endl;
            return false;
        }

        // 设置套接字选项
        int opt = 1;
        if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) == -1) {
            std::cerr << "Failed to set socket options" << std::endl;
            close(server_fd);
            return false;
        }

        // 绑定套接字
        sockaddr_in address;
        memset(&address, 0, sizeof(address));
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = INADDR_ANY;
        address.sin_port = htons(port);

        if (bind(server_fd, (struct sockaddr *)&address, sizeof(address)) == -1) {
            std::cerr << "Failed to bind socket" << std::endl;
            close(server_fd);
            return false;
        }

        // 监听连接
        if (listen(server_fd, 3) == -1) {
            std::cerr << "Failed to listen" << std::endl;
            close(server_fd);
            return false;
        }

        std::cout << "Server started on port " << port << std::endl;
        running = true;
        serverThread = std::thread(&BridgeServer::acceptConnections, this);
        return true;
    }

    void stop() {
        if (!running) return;
        running = false;
        
        // 关闭服务器套接字
        close(server_fd);
        
        // 关闭所有客户端连接
        {
            std::lock_guard<std::mutex> lock(clientMutex);
            for (int sock : clientSockets) {
                close(sock);
            }
            clientSockets.clear();
        }
        
        // 等待服务器线程结束
        if (serverThread.joinable()) {
            serverThread.join();
        }
        
        std::cout << "Server stopped" << std::endl;
    }

private:
    void acceptConnections() {
        while (running) {
            sockaddr_in clientAddress;
            socklen_t clientAddressSize = sizeof(clientAddress);
            
            // 接受客户端连接
            int clientSocket = accept(server_fd, (struct sockaddr *)&clientAddress, &clientAddressSize);
            if (clientSocket == -1) {
                if (running) {
                    std::cerr << "Failed to accept client connection" << std::endl;
                }
                continue;
            }

            // 打印客户端信息
            char clientIP[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, &clientAddress.sin_addr, clientIP, INET_ADDRSTRLEN);
            std::cout << "New client connected: " << clientIP << ":" << ntohs(clientAddress.sin_port) << std::endl;

            // 添加到客户端列表（线程安全）
            {
                std::lock_guard<std::mutex> lock(clientMutex);
                clientSockets.push_back(clientSocket);
            }

            // 为客户端创建单独的线程处理
            std::thread([this, clientSocket]() {
                handleClient(clientSocket);
                
                // 从列表中移除已断开的客户端（线程安全）
                std::lock_guard<std::mutex> lock(clientMutex);
                auto it = std::find(clientSockets.begin(), clientSockets.end(), clientSocket);
                if (it != clientSockets.end()) {
                    clientSockets.erase(it);
                }
            }).detach();
        }
    }

    void handleClient(int clientSocket) {
        char buffer[1024] = {0};
        
        while (running) {
            // 清空缓冲区
            memset(buffer, 0, sizeof(buffer));
            
            // 读取客户端消息
            int bytesRead = read(clientSocket, buffer, sizeof(buffer));
            
            if (bytesRead <= 0) {
                // 客户端断开连接
                std::cout << "Client disconnected" << std::endl;
                close(clientSocket);
                break;
            }
            
            // 处理接收到的消息
            std::string message(buffer, bytesRead);
            std::cout << "Received: " << message << std::endl;
            
            // 根据消息内容发送响应
            std::string response;
            if (message.find("START") != std::string::npos) {
                response = "ACK_START";
            } else if (message.find("STOP") != std::string::npos) {
                response = "ACK_STOP";
            } else if (message.find("heartbeat") != std::string::npos) {
                // 模拟JSON响应
                response = "{\"type\":\"heartbeat\",\"status\":\"alive\"}";
            } else {
                response = "UNKNOWN_COMMAND";
            }

            // response = R"({"type" : "heartbeat", "status" : "alive"})";
            // 发送响应
            ssize_t bytesSent = send(clientSocket, response.c_str(), response.length(), 0);
            if (bytesSent == -1) {
                std::cerr << "Failed to send response" << std::endl;
                close(clientSocket);
                break;
            }
            std::cout << "Sent: " << response << std::endl;
        }
    }
};

int main() {
    BridgeServer server(5002);
    
    // 启动服务器
    if (server.start()) {
        std::cout << "Server running. Press Enter to stop..." << std::endl;
        std::cin.get();
        server.stop();
    }
    
    return 0;
}